# -*- coding: utf-8 -*-

import sys
from PyQt5 import QtCore, QtGui
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QApplication, QDialog
from PyQt5.QtGui import QImage, QPixmap, QPainter, QPen
from od.label.od_label_ui import Ui_od_label

from od.label.od_label_backend import LabelBackend


class ODLabelWindow(QDialog, Ui_od_label):
    def __init__(self, bus_date_folder):
        super(ODLabelWindow, self).__init__()
        self.setupUi(self)
        self.backend = LabelBackend(bus_date_folder)
        
        self.pushButton_gallery_previous_person.clicked.connect(self.gallery_previous_person)
        self.pushButton_gallery_next_person.clicked.connect(self.gallery_next_person)
        self.pushButton_gallery_previous_image.clicked.connect(self.gallery_previous_image)
        self.pushButton_gallery_next_image.clicked.connect(self.gallery_next_image)
        self.pushButton_query_previous_person.clicked.connect(self.query_previous_person)
        self.pushButton_query_next_person.clicked.connect(self.query_next_person)
        self.pushButton_query_previous_image.clicked.connect(self.query_previous_image)
        self.pushButton_query_next_image.clicked.connect(self.query_next_image)
        self.pushButton_selected.clicked.connect(self.selected)
        self.pushButton_save.clicked.connect(self.save_and_go_on)

        self.label_gallery.setScaledContents(True)  # 图像铺满Label
        self.label_query.setScaledContents(True)
        self.label_selected_flag.setScaledContents(True)

        self.pushButton_save.setEnabled(False)

        self.show_image()

    @staticmethod
    def draw_rectangle(image, rectangle):
        painter = QPainter()
        painter.begin(image)
        painter.setPen(QPen(QtCore.Qt.red, 3))
        painter.drawRect(*rectangle)
        painter.end()

    def show_image(self):
        # gallery
        gallery_person = self.backend.gallery_person
        gallery_person_image_index = self.backend.gallery_person_image_index
        self.lineEdit_gallery_person_index.setText(
            F"{self.backend.gallery_persons_index+1}/{len(self.backend.gallery_persons)}"
        )
        self.lineEdit_gallery_image_index.setText(
            F"{gallery_person_image_index+1}/{len(gallery_person['image'])}"
        )
        self.lineEdit_gallery_time.setText(gallery_person['time'])
        self.lineEdit_gallery_location.setText(gallery_person['location'])

        gallery_image = QImage(gallery_person['image'][gallery_person_image_index])
        gallery_rectangle = gallery_person['rectangle'][gallery_person_image_index]
        self.draw_rectangle(gallery_image, gallery_rectangle)
        self.label_gallery.setPixmap(QPixmap(gallery_image))
        # query
        query_person = self.backend.query_person
        query_person_image_index = self.backend.query_person_image_index
        self.lineEdit_query_person_index.setText(
            F"{self.backend.query_persons_index+1}/{len(self.backend.query_persons)}"
        )
        self.lineEdit_query_image_index.setText(
            F"{query_person_image_index+1}/{len(query_person['image'])}"
        )
        self.lineEdit_query_time.setText(query_person['time'])
        self.lineEdit_query_location.setText(query_person['location'])

        query_image = QImage(query_person['image'][query_person_image_index])
        query_rectangle = query_person['rectangle'][query_person_image_index]
        self.draw_rectangle(query_image, query_rectangle)
        self.label_query.setPixmap(QPixmap(query_image))

        if self.pushButton_selected.text() == "取消":
            self.label_selected_flag.setPixmap(QPixmap("../images/right.png"))
        else:
            self.label_selected_flag.setPixmap(QPixmap("../images/question.png"))
        
    def gallery_previous_person(self):
        self.backend.gallery_persons_index -= 1
        self.show_image()
        
    def gallery_next_person(self):
        self.backend.gallery_persons_index += 1
        self.show_image()

    def gallery_previous_image(self):
        self.backend.gallery_person_image_index -= 1
        self.show_image()

    def gallery_next_image(self):
        self.backend.gallery_person_image_index += 1
        self.show_image()

    def query_previous_person(self):
        self.backend.query_persons_index -= 1
        self.show_image()

    def query_next_person(self):
        self.backend.query_persons_index += 1
        self.show_image()
        
    def query_previous_image(self):
        self.backend.query_person_image_index -= 1
        self.show_image()

    def query_next_image(self):
        self.backend.query_person_image_index += 1
        self.show_image()

    def selected(self):
        if self.pushButton_selected.text() == "取消":
            self.pushButton_selected.setText(r"确定是同一个人")
            self.pushButton_gallery_previous_person.setEnabled(True)
            self.pushButton_gallery_next_person.setEnabled(True)
            self.pushButton_query_previous_person.setEnabled(True)
            self.pushButton_query_next_person.setEnabled(True)
            self.pushButton_save.setEnabled(False)
        else:
            self.pushButton_selected.setText("取消")
            self.pushButton_gallery_previous_person.setEnabled(False)
            self.pushButton_gallery_next_person.setEnabled(False)
            self.pushButton_query_previous_person.setEnabled(False)
            self.pushButton_query_next_person.setEnabled(False)
            self.pushButton_save.setEnabled(True)
        self.show_image()

    def save_and_go_on(self):
        if self.pushButton_selected.text() != "取消":
            return
        self.backend.save_labeled_pair()
        self.selected()

    def event(self, event):
        if isinstance(event, QtGui.QKeyEvent) and event.type() == QtCore.QEvent.KeyPress:
            if event.key() == Qt.Key_W:
                self.gallery_previous_person()
            elif event.key() == Qt.Key_S:
                self.gallery_next_person()
            elif event.key() == Qt.Key_A:
                self.gallery_previous_image()
            elif event.key() == Qt.Key_D:
                self.gallery_next_image()
            elif event.key() == Qt.Key_I:
                self.query_previous_person()
            elif event.key() == Qt.Key_K:
                self.query_next_person()
            elif event.key() == Qt.Key_J:
                self.query_previous_image()
            elif event.key() == Qt.Key_L:
                self.query_next_image()
            else:
                pass
        return super().event(event)


def main():
    bus_date_folder = r"E:\data\59092_00E0B453CE0B\2019_3_18"
    app = QApplication(sys.argv)
    window = ODLabelWindow(bus_date_folder)
    window.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()
